import os
import cv2
import numpy as np
from pathlib import Path

def align_images(image1, image2):
    """
    使用SIFT特征点匹配对齐两张图片
    """
    # 转换为灰度图
    gray1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    
    # 初始化SIFT检测器
    sift = cv2.SIFT_create()
    
    # 检测关键点和描述符
    keypoints1, descriptors1 = sift.detectAndCompute(gray1, None)
    keypoints2, descriptors2 = sift.detectAndCompute(gray2, None)
    
    # 使用FLANN匹配器
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    matches = flann.knnMatch(descriptors1, descriptors2, k=2)
    
    # 筛选好的匹配点
    good_matches = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good_matches.append(m)
    
    # 获取匹配点的坐标
    src_pts = np.float32([keypoints1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    dst_pts = np.float32([keypoints2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    
    # 计算单应性矩阵
    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
    
    # 对齐图像
    aligned_image = cv2.warpPerspective(image1, M, (image2.shape[1], image2.shape[0]))
    
    return aligned_image

def combine_images(name_img, run_img):
    """
    拼接姓名图片的上部和跑步图片的下部
    """
    # 获取图片高度
    height = name_img.shape[0]
    
    # 计算分割位置 (2/9高度处)
    split_pos = int(height * 2 / 9)
    
    # 提取name图片的上部
    name_top = name_img[:split_pos, :]
    
    # 提取run图片的下部
    run_bottom = run_img[split_pos:, :]
    
    # 拼接两部分
    combined = np.vstack((name_top, run_bottom))
    
    return combined

def process_images(name_dir, run_dir, output_dir):
    """
    处理所有图片
    """
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 获取文件名列表
    name_files = sorted(os.listdir(name_dir))
    run_files = sorted(os.listdir(run_dir))
    
    # 确保图片数量相同
    if len(name_files) != len(run_files):
        print("警告: name和run文件夹中的图片数量不一致!")
        return
    
    # 处理每对图片
    for name_file, run_file in zip(name_files, run_files):
        # 读取图片
        name_path = os.path.join(name_dir, name_file)
        run_path = os.path.join(run_dir, run_file)
        
        name_img = cv2.imread(name_path)
        run_img = cv2.imread(run_path)
        
        if name_img is None or run_img is None:
            print(f"无法读取图片: {name_path} 或 {run_path}")
            continue
        
        # 对齐图片
        aligned_name = align_images(name_img, run_img)
        
        # 拼接图片
        combined = combine_images(aligned_name, run_img)
        
        # 保存结果
        output_path = os.path.join(output_dir, f"combined_{name_file}")
        cv2.imwrite(output_path, combined)
        print(f"已保存: {output_path}")

if __name__ == "__main__":
    # 设置路径
    name_dir = "nameSid"
    run_dir = "runPic"
    output_dir = "output"
    
    # 处理图片
    process_images(name_dir, run_dir, output_dir)
    
    print("所有图片处理完成!")